import torch
import torch.nn as nn
import math


class PositionalEmbedding(nn.Module):
    def __init__(self, d_model, max_seq_len=5000):
        super().__init__()
        # 定义记录每个token位置的索引，0-max_len-1
        # [max_len,1] 方便后续与缩放因子进行相处
        position = torch.arange(0, max_seq_len).unsqueeze(1)
        # 每个维度得到缩放因子， torch.arange(0, d_model, 2): 生成偶数维度索引 0 2 4 对应公式就是2i
        div_term = torch.exp(torch.arange(0, d_model, 2).float()*(-math.log(10000.0)/d_model))

        # d_model 每个词向量的维度； max_len:句子的最大长度
        # 初始化位置编码矩阵， 形状为max_len, d_model
        pe = torch.zeros(max_seq_len, d_model)

        # 每个token的位置索引pos*每个维度的缩放因子（div_term)再套上sin得到偶数维度的位置编码值
        pe[:, 0::2] = torch.sin(position*div_term)
        # 每个token的位置索引pos*每个维度的缩放因子（div_term)再套上cos得到奇数维度的位置编码值
        pe[:, 1::2] = torch.cos(position*div_term)
        # 增加batch维度， 1， max_len, d_model,方便后续与输入embedding进行相加
        pe = pe.unsqueeze(0)
        # 注册为buffer， 把位置编码pe存在模型里面，不参与训练，但是随着模型保存/加载
        self.register_buffer("pe", pe)

    def forward(self, x):
        # x: embedding 形状 batch, seq_len, d_model
        seq_len = x.shape[1]
        # 每个token的embedding加上对应位置的编码
        # self.pe[:,:seq_len,:] 取前seq_len个位置，形状会变成1，seq_len, d_model, 可以和输入X对齐
        # x+self.pe[:, :seq_len, :]:batch, seq_len，d_model；embedding加上位置编码，transformer就能知道token的顺序了
        return x+self.pe[:, :seq_len, :]


class SelfAttention(nn.Module):
    def __init__(self, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.softmax = nn.Softmax(dim=-1)  # 将得分转换为概率分布，在最后一个维度

    def forward(self, Q, K, V, mask=None):
        # X:batch, seq_len, d_model
        # d_model:embedding向量的维度
        # Q，query向量 维度：batch，heads，seq_len_q,d_k
        # k，key向量 维度：batch，heads，seq_len_k,d_k
        # v，value向量 维度：batch，heads，seq_len_v,d_v
        # mask:那些位置要看，哪些位置要忽略
        d_k = Q.size(-1)  # q的最后一维是对每一个query向量的维度，代表对每个query进行缩放
        # batch，heads，seq_len_q,d_k,   batch，heads,d_k,seq_len_k ---> batch，heads,seq_len_q,seq_len_k
        score = torch.matmul(Q, K.transpose(-2, -1))/math.sqrt(d_k)  # 进行缩放，让梯度更稳定
        # 如果提供了mask，则通过mask==0来找到需要屏蔽的位置，mask_fill会将这些为宗旨的值改为-inf(负无穷)
        # 然后经过softmax之后值会变为0
        # 设置mask==0 表示被屏蔽 mask==1表示当前位置可见
        if mask is not None:
            # print("mask shape:", mask.shape)
            score = score.masked_fill(mask == 0, float('-inf'))
        # batch，heads,seq_len_q,seq_len_k 对最后一维进行softmax，即对key进行，得到注意力权重矩阵，对每一个query的key权重之和为1
        att = self.softmax(score)
        att = self.dropout(att)  # 对注意力权重进行dropout，防止过拟合
        # att:batch，heads,seq_len_q,seq_len_k; V:batch，heads，seq_len_v,d_v ---> att*V: batch，heads，seq_len_q, d_v
        out = torch.matmul(att, V)
        return out, att


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_head, dropout=0.1):
        super().__init__()
        # d_model: embedding的维度 512
        # n_head: 多头注意力的头数 8
        # d_model需要被 n_head整除 结果为64
        assert d_model % n_head == 0
        self.d_k = d_model // n_head  # 每个头的维度
        self.n_head = n_head

        # 将输入映射到Q， K， V三个向量，通过线性映射让模型具有学习能力
        self.W_q = nn.Linear(d_model, d_model)  # query的线性映射，维度不需要改变，方便后续的多头拆分
        self.W_k = nn.Linear(d_model, d_model)  # key的线性映射
        self.W_v = nn.Linear(d_model, d_model)  # value的线性映射
        self.fc = nn.Linear(d_model, d_model)   # 多头拼接后再映射回原来的d_model，让模型融合不同头的信息

        self.attention = SelfAttention(dropout)
        self.dropout = nn.Dropout(dropout)
        self.norm = nn.LayerNorm(d_model)   # 用于残差后的归一化

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)  # 获取batch的大小
        # q 的维度 batch，seq_len,d_model -->batch,seq_len,self.n_head,self.d_k --> batch,self.n_head,seq_len,self.d_k
        # 为了让每个注意力头独立处理整个序列，方便后续计算注意力权重
        Q = self.W_q(q).view(batch_size, -1, self.n_head, self.d_k).transpose(1, 2)
        K = self.W_k(k).view(batch_size, -1, self.n_head, self.d_k).transpose(1, 2)
        V = self.W_v(v).view(batch_size, -1, self.n_head, self.d_k).transpose(1, 2)

        # 计算注意力
        out, att = self.attention(Q, K, V, mask)    # att为注意力权重， out为注意力加权后的值
        # out.transpose(1, 2): batch，heads，seq_len_q, d_v --> batch，seq_len_q, heads，d_v --> batch,seq_len,d_model
        # contiguous目的是让tensor在内存中连续存储，避免view的时候产生报错
        # 多头拼接
        out = out.transpose(1, 2).contiguous().view(batch_size, -1, self.n_head*self.d_k)
        out = self.fc(out)  # 让输入输出一致，方便残差连接
        out = self.dropout(out)
        # 残差连接+layernorm
        return self.norm(out+q), att  # 返回输出和注意力权重
        # return self.norm(out+q)  # 返回输出和注意力权重


class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        # d_ff: 前馈隐藏层维度
        self.fc1 = nn.Linear(d_model, d_ff)  # 输入维度为d_model,输出为d_ff，512 ->2048， 为了让模型学到更丰富的特征
        self.fc2 = nn.Linear(d_ff, d_model)  # 保证第二个线性层输出维度等于第一个线性层的输入维度，为了后续做残差连接
        self.dropout = nn.Dropout(dropout)
        self.norm = nn.LayerNorm(d_model)  # 对最后一维进行归一化

    def forward(self, x):
        # x 的形状为 batch， seq_len, d_model
        out = self.fc2(self.dropout(torch.relu(self.fc1(x))))
        return self.norm(out + x)  # 残差为了保留输入的低阶信息，避免训练时信息丢失
        # 先残差再归一化，为了让模型更稳定，加快收敛


class EncoderLayer(nn.Module):
    def __init__(self, heads, d_model, d_ff, dropout=0.1):
        super().__init__()
        # 输入为src 实现序列内部的信息交互，每个token都可以看到序列中的其他token，从而学习到上下文依赖
        self.self_multi_head_att = MultiHeadAttention(d_model=d_model, n_head=heads, dropout=dropout)
        # 对每个位置向量独立进行非线性变换，提升模型表达能力
        self.ffn = FeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout)

    def forward(self, src, mask=None):
        # src 输入序列张量， 形状 batch， seq_len,d_model
        # mask屏蔽padding位置， 避免模型关注无效token
        # Q, K, V = src 对输入序列本身进行自注意力计算
        out, _ = self.self_multi_head_att(src, src, src, mask)
        # 经过前馈神经网络，每个位置的token都会单独通过两层线性层映射和激活函数，提升模型的表达能力
        out = self.ffn(out)
        # 返回编码后的结果
        return out


class DecoderLayer(nn.Module):
    def __init__(self, heads, d_model, d_ff, dropout=0.1):
        super().__init__()
        # Mask多头注意力机制
        # 输入tgt(目标序列) 在翻译任务中 已经生成的前几个词
        # 计算目标序列内部的自注意力，通过mask遮挡未来的token
        self.masked_att = MultiHeadAttention(d_model=d_model, n_head=heads, dropout=dropout)
        # 交叉注意力，和encoder做交互
        # 输入Q=当前解码器的输出，K＝V＝　来自编码器的memory（原序列上下文信息）
        # 为了将目标序列与原序列进行对齐
        self.att = MultiHeadAttention(d_model=d_model, n_head=heads, dropout=dropout)
        self.ffn = FeedForward(d_model=d_model, d_ff=d_ff, dropout=dropout)  # 为了提升模型的表达能力

    def forward(self, tgt, memory, tgt_mask=None, memory_mask=None):
        # tgt 目标序列 memory：编码器输出（原序列的表示）
        # tgt_mask:屏蔽未来的token， memory_mask: PAD做掩码

        # 目标序列内部的自注意力， 未来位置被mask
        out, _ = self.masked_att(tgt, tgt, tgt, tgt_mask)
        # 将目标序列和原序列进行交互，Q解码器当前的输出out，K=V=memory（编码器的输出）
        out, _ = self.att(out, memory, memory, memory_mask)
        out = self.ffn(out)
        return out


class Encoder(nn.Module):
    def __init__(self, vocab_size, d_model, num_layer, heads, d_ff, dropout=0.1, max_seq_len=5000):
        super().__init__()
        #
        self.embedding = nn.Embedding(vocab_size, d_model)
        # 位置编码加入序列中token的位置信息
        self.position_encode = PositionalEmbedding(d_model, max_seq_len)
        # 构建编码器的堆叠结构
        # 生成num_layer个encoder
        self.encode_layer = nn.ModuleList([EncoderLayer(heads, d_model, d_ff, dropout) for _ in range(num_layer)])

    def forward(self, x, src_mask):
        # 将输入tokenID转换为embedding向量
        # 输出　shape batch, seq_len, d_model
        # 乘上　sqrt(d_model), 进行缩放，　让后续注意力计算更稳定
        embed_x = self.embedding(x) * math.sqrt(self.embedding.embedding_dim)
        pos_encode_x = self.position_encode(embed_x)
        for layer in self.encode_layer:
            pos_encode_x = layer(pos_encode_x, src_mask)
        return pos_encode_x


class Decoder(nn.Module):
    def __init__(self, vocab_size, d_model, num_layer, heads, d_ff, dropout=0.1, max_seq_len=5000):
        super().__init__()
        # 将目标序列的tokenID转换为向量　维度为d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.position_encode = PositionalEmbedding(d_model, max_seq_len)

        self.decode_layer = nn.ModuleList([DecoderLayer(heads, d_model, d_ff, dropout) for _ in range(num_layer)])
        self.fc_out = nn.Linear(d_model, vocab_size)

    def forward(self, x, encoder_kv, dst_mask=None, src_dst_mask=None):
        embed_x = self.embedding(x) * math.sqrt(self.embedding.embedding_dim)
        pos_encode_x = self.position_encode(embed_x)
        for layer in self.decode_layer:
            pos_encode_x = layer(pos_encode_x, encoder_kv, dst_mask, src_dst_mask)
        # 将解码器最后一层输出的隐藏向量映射回原词汇表的维度，得到每个token的预测向量
        return self.fc_out(pos_encode_x)


class Transformer(nn.Module):
    def __init__(self, encode_vocab_size, decode_vocab_size, d_model=512, num_layer=6, heads=8, d_ff=2048, dropout=0.1, max_seq_len=5000):
        super().__init__()
        # 编码器 将目标语言token编码为上下文表示
        self.encoder = Encoder(encode_vocab_size, d_model, num_layer, heads, d_ff, dropout, max_seq_len)
        # 解码器 根据编码器的输出和目标语言输入生成预测
        self.decoder = Decoder(decode_vocab_size, d_model, num_layer, heads, d_ff, dropout, max_seq_len)

    def forward(self, src, tgt, src_mask=None, tgt_mask=None,memory_mask=None):
        memory = self.encoder(src, src_mask)
        out = self.decoder(tgt, memory, tgt_mask, memory_mask)
        # 返回transformer输出 batch， seq_len——tgt, tgt_vocab
        return out


def generate(size):
     mask =torch.triu(torch.ones(size, size), diagonal=1).bool()

     return mask == 0


if __name__ == "__main__":
    src = 10000
    tgt = 10000
    att = Transformer(src, tgt)
    x = torch.randint(0, 10000, (32, 10))
    y = torch.randint(0, 10000, (32, 20))
    tgt_mask = generate(y.size(1)).to(y.device)
    out = att(x, y, tgt_mask=tgt_mask)
    print(out.shape)









